
#include "aerBVH.hpp"


#include <cfloat>
#include <algorithm>
#include "../aerSortIndex.hpp"

#include "../aerOpenGL.hpp"// temp


namespace aer {


/* PUBLIC __________________ */

BVH::BVH()
    : m_numNodes(0u), 
      m_numLeaves(0u), 
      m_pMeshes(NULL),
      m_pBoundingBoxes(NULL)
{
  m_tree.clear();
  m_attribs.clear();
}

//void BVH::build(const size_t nElems, const BoundingBox *pBoundingBoxes)
void BVH::build(const size_t nElems, const Mesh *pMeshes, const BoundingBox *pBoundingBoxes)
{
  m_numLeaves = nElems;
  m_numNodes = 2u*m_numLeaves - 1u;
  m_pMeshes = pMeshes;//
  m_pBoundingBoxes = pBoundingBoxes;//
  
  m_tree.resize(m_numNodes);
  m_attribs.resize(m_numNodes);  
  
  
  init();//
  
  
  /// Sort objects' AABBs per axis
  std::vector<unsigned int> indicesX(m_numLeaves);
  std::vector<unsigned int> indicesY(m_numLeaves);
  std::vector<unsigned int> indicesZ(m_numLeaves);
  {
    const unsigned int nAttribs = 3u;
    const unsigned int nDatas = nAttribs * m_numLeaves;
    
    float datas[nDatas];
    
    for (unsigned int i=0u, j=0u; i<m_numLeaves; ++i, j+=nAttribs) 
    {
      const aer::Vector3 &position = m_pBoundingBoxes[i].getCenter();
      
      datas[j+0u] = position.x;
      datas[j+1u] = position.y;
      datas[j+2u] = position.z;
    }
    
    aer::SortIndex sortIndex;
    sortIndex.sortX( m_numLeaves, datas, indicesX.data());
    sortIndex.sortY( m_numLeaves, datas, indicesY.data());
    sortIndex.sortZ( m_numLeaves, datas, indicesZ.data());
  }
  
  /// LUT used to resort indices per-axis
  std::vector<unsigned int> resortLUT(m_numLeaves);
  
  /// Maps common buffers used during the build
  BuildParams params(indicesX.data(), indicesY.data(), indicesZ.data(), 
                     resortLUT.data());
  
  
  unsigned int idCounter = 0u;
  subBuild( 0u, 0u, &idCounter, params, 0u, m_numLeaves);
}



//--------------- XXX

void BVH::display(glm::vec4 planes[Frustum::kNumPlane], BVH::Node *node, unsigned int depth)
{
  if (node == 0) {
    display(planes, &m_tree[0u]);
    return;
  }

  // Frustum culling
  if (!isInsideViewFrustum( node->id, planes)) {
   return;
  }
  
  if (!IsALeaf( *node ))
  {
    unsigned int off = GetNodeOffset(*node);
    display( planes, &m_tree[off], depth+1u);
    display( planes, &m_tree[off+1u], depth+1u);

    const float c = depth / log2f( m_numLeaves );
    glColor3f( 0.2f, 0.45f*c+.2f, 0.9f*c + .01f );
  }
  else
  {
    glColor3f( 0.95f, 0.45f, 0.31f );
  }
  
  glPushMatrix();
  
  const Vector3 &pos   = m_attribs[node->id].aabb.getCenter();
  const Vector3 &scale = m_attribs[node->id].aabb.getScale();
  
  glTranslatef( pos.x, pos.y, pos.z );
  glScalef( scale.x, scale.y, scale.z );
  
  glBegin( GL_QUADS );
    glVertex3f(-1.0f, -1.0f, -1.0f);
    glVertex3f(-1.0f, -1.0f, +1.0f);
    glVertex3f(-1.0f, +1.0f, +1.0f);
    glVertex3f(-1.0f, +1.0f, -1.0f);
    
    glVertex3f(+1.0f, -1.0f, -1.0f);
    glVertex3f(+1.0f, -1.0f, +1.0f);
    glVertex3f(+1.0f, +1.0f, +1.0f);
    glVertex3f(+1.0f, +1.0f, -1.0f);
    
    glVertex3f(-1.0f, +1.0f, -1.0f);
    glVertex3f(-1.0f, +1.0f, +1.0f);
    glVertex3f(+1.0f, +1.0f, +1.0f);
    glVertex3f(+1.0f, +1.0f, -1.0f);
    
    glVertex3f(-1.0f, -1.0f, -1.0f);
    glVertex3f(-1.0f, -1.0f, +1.0f);
    glVertex3f(+1.0f, -1.0f, +1.0f);
    glVertex3f(+1.0f, -1.0f, -1.0f);
  glEnd(); 
  glPopMatrix();
}
//--------------- XXX



/* PRIVATE __________________ */


void BVH::subBuild(const unsigned int nodeId, const unsigned int parentId, unsigned int *pIdCounter,
                   const BuildParams &params, const size_t offset, const size_t count)
{
  Node &node = m_tree[nodeId];
  
  /// Compute node BBOX 
  BoundingBox &bbox = m_attribs[nodeId].aabb; //
  {
    unsigned int *indices = &params.pIndices[0][offset];
    
    bbox = BoundingBox( m_pBoundingBoxes[ indices[0u] ] );
    
    for (unsigned int i=1u; i<count; ++i) {
      bbox.addAABB( m_pBoundingBoxes[ indices[i] ] );
    }
  }
  
  
  /// Is a Leaf ? return !
  if (count <= 1)
  {
    unsigned int aabbId = params.pIndices[0][offset];
    node.set( nodeId, parentId, aabbId, true);
    return;
  }
  
  
  /// Find Best cost for the Surface Area Heuristic (along the 3 axis)
  float         bestSA_cost = FLT_MAX;
  unsigned int  bestSA_id   = 0u;
  unsigned char bestSA_axis = 0;
    
  for (unsigned char axis = 0; axis < 3; ++axis)
  {
    if (findBestSACost( params, offset, count, axis, &bestSA_cost, &bestSA_id)) {
      bestSA_axis = axis;
    }
  }
  
  /// Resort others axis (left & right subset)
  if (bestSA_axis != 0) resort( params, offset, count, bestSA_axis, 0, bestSA_id);
  if (bestSA_axis != 1) resort( params, offset, count, bestSA_axis, 1, bestSA_id);
  if (bestSA_axis != 2) resort( params, offset, count, bestSA_axis, 2, bestSA_id);
  
  
  /// Continue recursively  
  size_t leftOffset  = offset;
  size_t leftCount   = bestSA_id;
  size_t rightOffset = leftOffset + leftCount;
  size_t rightCount  = count - leftCount;
  
  
  unsigned leftId  = (*pIdCounter) + 1u;
  unsigned rightId = (*pIdCounter) + 2u;
  *pIdCounter = rightId;
  
  node.set( nodeId, parentId, leftId, false);  

  subBuild( leftId,  nodeId, pIdCounter, params, leftOffset,  leftCount);
  subBuild( rightId, nodeId, pIdCounter, params, rightOffset, rightCount);
}


bool BVH::findBestSACost(const BuildParams &params, 
                         const size_t offset, 
                         const size_t count, 
                         const unsigned char axis, 
                         float *pBestCost, 
                         unsigned int *pBestId)
{
  bool bNewBest = false;
  
  const unsigned int *indices = &params.pIndices[axis][offset];

  /// Left bounding box
  BoundingBox leftBox( m_pBoundingBoxes[ indices[0u] ] );

  /// Right bounding boxes stored 'right to left'
  const size_t nRightBoxes = count-1u;
  
  std::vector<BoundingBox> rightBox(nRightBoxes); // put outside ?  
  rightBox[0u] = m_pBoundingBoxes[ indices[nRightBoxes] ];

  /// Right boxes are precomputed to be in linear time
  for (unsigned int i=1u; i<nRightBoxes; ++i)
  {
    rightBox[i] = BoundingBox( rightBox[i-1u] );
    
    const size_t idx = indices[nRightBoxes - i];
    rightBox[i].addAABB( m_pBoundingBoxes[idx] );
  }  
  
  // Unused current box surface area to ponderate the cost
  //const float invSA = 1.0f / bbox.getSurfaceArea();
  
  /// Search minimal cost
  for (unsigned int i=1u; i<count; ++i)
  {
    float lSA = leftBox.getSurfaceArea();
    float rSA = rightBox[nRightBoxes - i].getSurfaceArea();
    
    // Compute Surface Area cost
    float cost = i * lSA + (count-i) * rSA; //   + glm::max(lSA,rSA)
  
    if (cost < *pBestCost) 
    {
      *pBestCost = cost;
      *pBestId   = i;
      bNewBest   = true;
    }
    
    // update left Bbox
    leftBox.addAABB( m_pBoundingBoxes[ indices[i] ] );
  }
  
  return bNewBest;
}


/// resorts per-axis datas in linear time
void BVH::resort( const BuildParams &params,
                  const size_t offset, 
                  const size_t count,
                  const unsigned char baseAxis, 
                  const unsigned char axis, 
                  const unsigned int splitIdx )
{  
  unsigned int *baseIndices = &params.pIndices[baseAxis][offset];
  unsigned int *axisIndices = &params.pIndices[axis][offset];
 
  
  /// Fill the LUT (partially)
  unsigned int *LUT = params.pLUT;  
  for (unsigned int i=0u; i<count; ++i) {
    LUT[ baseIndices[i] ] = i;
  }  
  
  /// Resort left & right subset
  unsigned int leftIdx = 0u;
  unsigned int rightIdx = splitIdx;
  
  // tmp buffer to store resorted values
  unsigned int resortedBuffer[ count ];//
  
  /// resort in linear time
  for (unsigned int i=0u; i<count; ++i)
  {
    unsigned int idx = axisIndices[i];
    
    if (LUT[idx] < splitIdx) {
      resortedBuffer[leftIdx++] = idx;
    } else {
      resortedBuffer[rightIdx++] = idx;
    }
  }
  
  std::copy( resortedBuffer, &resortedBuffer[count], axisIndices);  
}

} // aer
